import pandas as pd
import splunklib.client as client
from splunklib import results


def get_dangerous_ips(config_yaml, filter):
    service = client.connect(
        host=config_yaml['splunk']['host'],
        port=config_yaml['splunk']['port'],
        scheme=config_yaml['splunk']['scheme'],
        username=config_yaml['splunk']['username'],
        password=config_yaml['splunk']['password']
    )
    start_time = config_yaml['session']["start_time"]
    end_time = config_yaml['session']["end_time"]
    job = service.jobs.oneshot(
        "search index=idx_pa" +
        filter + """
THREAT AND NOT informational AND  sourcetype="pan:threat"
| rex field=_raw  "THREAT,(?P<LOG_TYPE>.+?),.*?,(?P<PA_DATE>.*?),(?P<SIP>.*?),(?P<DIP>.*?),(?:.*?,.*?,.*?){7},(?P<S_PORT>.*?),(?P<D_PORT>.*?),.*?,.*?,.*?,(?P<PROTOCOL>.*?),(?P<DENY_METHOD>.*?),(?P<THREAT_SUMMARY>.*?),(?P<SEVERITY>medium|high|critical|low),"
| eval RAW_BAK=_raw
| eval THREAT_TIME = strftime(strptime(PA_DATE, "%Y/%m/%d %H:%M:%S"), "%Y-%m-%d %H:%M:%S")
| rex mode=sed field=_raw "s/.*?,\w{8}-\w{4}-\w{4}-\w{4}-\w{12},/start_ama_flag,/g"
| rex field=_raw "start_ama_flag,.*?,.*?,(?<XFF_IP>.*?),"
| eval _raw=RAW_BAK
| table THREAT_TIME,SIP,S_PORT, DIP, D_PORT,XFF_IP PROTOCOL, DENY_METHOD, THREAT_SUMMARY, SEVERITY
| dedup THREAT_TIME,SIP,S_PORT, DIP, D_PORT,XFF_IP PROTOCOL
    """, **{
            "earliest_time": start_time.strftime('%Y-%m-%dT%H:%M:%S'),
            "latest_time": end_time.strftime('%Y-%m-%dT%H:%M:%S'),
            "output_mode": "json",
            "count": 20000
        })
    dangerous_list = [item for item in results.JSONResultsReader(job) if isinstance(item, dict)]
    return pd.DataFrame(dangerous_list)
